
#include "fun.h"
#include "hs_cuda_outputwise.cuh"

struct FillValueFunctor
{

    FillValueFunctor(float fill_value)
        : m_fill_value(fill_value)
    { } 

    __device__ void operator()(float & out_val)
    {
        out_val = m_fill_value;
    }

private:
    float m_fill_value;

};

void rasterFillValue(hs::raster::RasterData<float> & mat, float fill_value)
{
    /* 将输入数据拷贝到GPU */
    hscuda::raster::RasterData<float> device_mat;
    hscuda::raster::host2device(mat.view(), device_mat);

    /* 调用 outputwise 模板函数 */
    int N = mat.view().numel();
    hs::cuda::outputwise::Unary(N
        , FillValueFunctor(fill_value)
        , device_mat.view().pData());
    cudaDeviceSynchronize();

    /* 将数据拷贝回主机设备 */
    device_mat.toHost(mat);
}


struct SubRasterFunctor
{
    SubRasterFunctor(int offset_x, int offset_y, int sub_size_x, int sub_size_y, int src_size_x, int src_size_y, float* p_src_datas)
        : m_offset_x(offset_x), m_offset_y(offset_y), m_sub_size_x(sub_size_x), m_sub_size_y(sub_size_y), m_src_size_x(src_size_x), m_src_size_y(src_size_y), m_p_src_datas(p_src_datas)
    { }

    __device__ void operator()(int id, float & out_val)
    {
        if (id < m_sub_size_x * m_sub_size_y) {
            // 将输出数据的线性索引转换为平面索引
            int sub_yid = id / m_sub_size_x;
            int sub_xid = id - sub_yid * m_sub_size_x;

            sub_yid += m_offset_y;
            sub_xid += m_offset_x;

            // 将平面索引转换为输入原图像的线性索引
            int src_ind = sub_xid + sub_yid * m_src_size_x;
            if (src_ind < m_src_size_x * m_src_size_y) {
                out_val = m_p_src_datas[src_ind];
            }
            else {
                out_val = 0;
            }
        }
    }

private:
    int m_offset_x;
    int m_offset_y;
    int m_sub_size_x;
    int m_sub_size_y;

    int m_src_size_x;
    int m_src_size_y;
    float* m_p_src_datas; 
};

void getSubRaster(const hs::raster::RasterData<float> & src_mat, int offset_x, int offset_y, int sub_size_x, int sub_size_y, hs::raster::RasterData<float> & dst_mat)
{
    // 将数据拷贝到GPU端
    hscuda::raster::RasterData<float> device_src_mat;
    hscuda::raster::host2device(src_mat.view(), device_src_mat);

    // 为输出分配显存
    hscuda::raster::RasterData<float> device_dst_mat(sub_size_x, sub_size_y);

    // 调用 outputwiseIndex
    int src_size_x = src_mat.view().sizeX();
    int src_size_y = src_mat.view().sizeY();
    int N = device_dst_mat.view().numel();
    hs::cuda::outputwiseIndex::Unary(N
        , SubRasterFunctor(offset_x, offset_y
            , sub_size_x, sub_size_y
            , src_size_x, src_size_y
            , device_src_mat.view().pData())
        , device_dst_mat.view().pData()
    );
    cudaDeviceSynchronize();

    // 将结果拷贝到主机端
    device_dst_mat.toHost(dst_mat);
}

